
import sys
sys.path.append("/media/cfs/liuhongru3/medicine/bert_seq2seq-master/")
sys.path.append("/media/cfs/liuhongru3/medicine/")

import torch 
from tqdm import tqdm
import torch.nn as nn 
from torch.optim import Adam
import pandas as pd
import numpy as np
import os
import json
import time
from torch.utils.data import Dataset, DataLoader
from bert_seq2seq.tokenizer import Tokenizer, load_chinese_base_vocab
from bert_seq2seq.utils import load_bert, load_model_params, load_recent_model
import opencc 
from rouge import Rouge

max_p_len = 200#128#512
max_q_len = 64#28
max_a_len = 128#150
def text_segmentate(text, maxlen, seps='\n', strips=None):
    """将文本按照标点符号划分为若干个短句
    """
    text = text.strip().strip(strips)
    if seps and len(text) > maxlen:
        pieces = text.split(seps[0])
        text, texts = '', []
        for i, p in enumerate(pieces):
            if text and p and len(text) + len(p) > maxlen - 1:
                texts.extend(text_segmentate(text, maxlen, seps[1:], strips))
                text = ''
            if i + 1 == len(pieces):
                text = text + p
            else:
                text = text + p + seps[0]
        if text:
            texts.extend(text_segmentate(text, maxlen, seps[1:], strips))
        return texts
    else:
        return [text]

def read_corpus(data_path, vocab_path):
    """
    读原始数据
    """
    word2idx = load_chinese_base_vocab(vocab_path, simplfied=True)
    tokenizer = Tokenizer(word2idx)
    trainval_data = json.load(open(data_path))
    seps, strips = u'\n。！？!?；;，, ', u'；;，, '
    data =[]
    for d in trainval_data:
        for p in d['annotations']:
            if p['A']:
                #data.append((d['text'],p['Q'],p['A']))
                for t in text_segmentate(d['text'], 200 - 2, seps, strips):
                    if p['A'] in t:
                        ## 加入数据之前过滤下一些特殊字符
                        t = t.replace("”", "").replace("“", "").replace("——", "，").replace("；", "，")
                        data.append((t, p['Q'], p['A']))
    # 保存一个随机序（供划分valid用）
    if not os.path.exists('./random_order.json'):
        random_order = list(range(len(data)))
        np.random.shuffle(random_order)
        json.dump(random_order, open('./random_order.json', 'w'), indent=4)
    else:
        random_order = json.load(open('./random_order.json'))
    # train_data = [data[j] for i, j in enumerate(random_order) if i % 10 != 0]
    train_data = [data[j] for i, j in enumerate(random_order)]
    valid_data = [data[j] for i, j in enumerate(random_order) if i % 10 == 0]
    return train_data,valid_data

class BertDataset(Dataset):
    """
    针对特定数据集，定义一个相关的取数据的方式
    """
    def __init__(self, passage,Q, A, vocab_path) :
        ## 一般init函数是加载所有数据
        super(BertDataset, self).__init__()
        # 读原始数据
        # self.sents_src, self.sents_tgt = read_corpus(poem_corpus_dir)
        self.passage = passage
        self.Q = Q
        self.A = A
        self.word2idx = load_chinese_base_vocab(vocab_path, simplfied=True)
        self.idx2word = {k: v for v, k in self.word2idx.items()}
        self.tokenizer = Tokenizer(self.word2idx)

    def __getitem__(self, i):
        ## 得到单个数据
        passage = self.passage[i]
        Q = self.Q[i]
        A = self.A[i]
        # print(pa)
        p_token_ids, p_token_type_ids = self.tokenizer.encode(passage, max_length=max_p_len)
        a_token_ids, a_token_type_ids = self.tokenizer.encode(A,max_length=max_a_len)
        q_token_ids, q_token_type_ids = self.tokenizer.encode(Q,max_length=max_q_len)
        token_ids = p_token_ids + a_token_ids[1:] + q_token_ids[1:]
        token_type_ids = [0]*(len(p_token_ids)+len(a_token_ids[1:]))
        token_type_ids += [1]*(len(q_token_ids[1:]))
        #print('token_ids',token_ids,'token_type_ids',token_type_ids)
        output = {
            "token_ids": token_ids,
            "token_type_ids": token_type_ids,
        }
        return output

    def __len__(self):

        return len(self.Q)

def collate_fn(batch):
    """
    动态padding， batch为一部分sample
    """

    def padding(indice, max_length, pad_idx=0):
        """
        pad 函数
        """
        pad_indice = [item + [pad_idx] * max(0, max_length - len(item)) for item in indice]
        return torch.tensor(pad_indice)

    token_ids = [data["token_ids"] for data in batch]
    max_length = max([len(t) for t in token_ids])
    token_type_ids = [data["token_type_ids"] for data in batch]

    token_ids_padded = padding(token_ids, max_length)
    token_type_ids_padded = padding(token_type_ids, max_length)
    target_ids_padded = token_ids_padded[:, 1:].contiguous()

    return token_ids_padded, token_type_ids_padded, target_ids_padded

class QuestionTrainer:
    def __init__(self):
        # 加载数据
        self.vocab_path = "./roberta_wwm_vocab.txt" # roberta模型字典的位置
        # self.vocab_path = "./vocab.txt"
        #self.sents_src, self.sents_tgt = read_corpus(data_dir + "/Poetry1", self.vocab_path)
        self.train_data,self.valid_data = read_corpus("./round1_train_0907.json", self.vocab_path)
        self.model_name = "roberta" # 选择模型名字
        self.model_path = "./roberta_wwm_pytorch_model.bin" # roberta模型位置
        # self.model_path = "../pytorch_model.bin" # roberta模型位置
        self.recent_model_path = "./bert_model_medicine.bin" # 用于把已经训练好的模型继续训练
        self.model_save_path = "./bert_model_medicine.bin"
        self.batch_size = 8
        self.lr = 1e-5
        # 判断是否有可用GPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("device: " + str(self.device))
        # 定义模型
        self.bert_model = load_bert(self.vocab_path, model_name=self.model_name, simplfied=True)
        
        #self.bert_model.load_state_dict(torch.load(self.model_save_path))
        
        ## 加载预训练的模型参数～
        load_model_params(self.bert_model, self.model_path)
        # 将模型发送到计算设备(GPU或CPU)
        self.bert_model.to(self.device)
        # 声明需要优化的参数
        self.optim_parameters = list(self.bert_model.parameters())
        self.optimizer = torch.optim.Adam(self.optim_parameters, lr=self.lr, weight_decay=1e-5)
        # 声明自定义的数据加载器
        passage,Q,A =[x[0] for x in self.train_data],[x[1] for x in self.train_data],[x[2] for x in self.train_data]
        dataset = BertDataset(passage, Q, A, self.vocab_path)
        self.dataloader =  DataLoader(dataset, batch_size=self.batch_size, shuffle=True, collate_fn=collate_fn)

    def train(self, epoch):
        # 一个epoch的训练
        self.bert_model.train()
        self.iteration(epoch, dataloader=self.dataloader, train=True)
    
    def save(self, save_path):
        """
        保存模型
        """
        torch.save(self.bert_model.state_dict(), save_path)
        print("{} saved!".format(save_path))
        
    def load(self):
        self.bert_model.load_state_dict(torch.load(self.model_save_path))
        self.bert_model.eval()
        data = json.load(open('../round1_test_0907.json'))
        seps, strips = u'\n。！？!?；;，, ', u'；;，, '
        with open('qa.json', 'w', encoding='utf-8') as f:
            for index,d in enumerate(data):
                for i, p in enumerate(d['annotations']):
                    found =False
                    for t in text_segmentate(d['text'], max_p_len - 2, seps, strips):
                        if p['A'] in t and not found:
                            found =True
                            q = self.bert_model.generate(t,p['A'],beam_size=5,device=self.device)
                            q = q.replace('[CLS]','')
                            p['Q'] =q
                            break
                    if not found:
                        q = self.bert_model.generate(d['text'],p['A'],beam_size=5,device=self.device)
                        q = q.replace('[CLS]','')
                        p['Q'] =q
                    print('A',p['A'],'Q',p['Q'])
                    data[index]['annotations'][i] = p
            json.dump(data,f,ensure_ascii=False)
                
        
    def iteration(self, epoch, dataloader, train=True):
        total_loss = 0
        start_time = time.time() ## 得到当前时间
        step = 0
        report_loss = 0
        # for token_ids, token_type_ids, target_ids in tqdm(dataloader,position=0, leave=True):
        for token_ids, token_type_ids, target_ids in tqdm(dataloader, position=0, leave=True):
            #print('token_ids',token_ids,'token_type_ids',token_type_ids,'target_ids',target_ids)
            step += 1
            token_ids = token_ids.to(self.device)
            token_type_ids = token_type_ids.to(self.device)
            target_ids = target_ids.to(self.device)
            # 因为传入了target标签，因此会计算loss并且返回
            predictions, loss = self.bert_model(token_ids,
                                                token_type_ids,
                                                labels=target_ids,
                                                device=self.device
                                                )
            
            # 反向传播
            if train:
                # 清空之前的梯度
                self.optimizer.zero_grad()
                # 反向传播, 获取新的梯度
                loss.backward()
                # 用获取的梯度更新模型参数
                self.optimizer.step()

            # 为计算当前epoch的平均loss
            total_loss += loss.item()
            report_loss += loss.item()

            if step % 500 == 0:
                # 看看输出什么情况
                passage = "橄榄，又名青果、白榄，为橄榄科植物橄榄的果实，产广东、广西、福建等地。宋朝大文学家苏东坡称之为“青子”。早在唐宋之间，橄榄已广泛地被采入药用。现代研究证实，橄榄的果实中含有蛋白质、脂肪、碳水化合物以及钙.磷、铁等。"
                a = "橄榄的果实中含有蛋白质、脂肪、碳水化合物以及钙.磷、铁等。"
                print(self.bert_model.muti_generete(passage, a, beam_size=3, device=self.device))
                print("loss is :" + str(report_loss))
                report_loss = 0

        end_time = time.time()
        spend_time = end_time - start_time
        # 打印训练信息
        print("epoch is " + str(epoch)+". loss is " + str(total_loss) + ". step is" + str(step) + ". spend time is "+ str(spend_time))
        report_loss = 0
        self.bert_model.eval()
        seps, strips = u'\n。！？!?；;，, ', u'；;，, '
        myRouge = Rouge()
        total_score,cnt =0.0,0
        for passage,Q,A in self.valid_data:
            found = False
            score = 0.0
            for t in text_segmentate(passage, max_p_len - 2, seps, strips):
                if A in t and not found:
                    found =True
                    res = self.bert_model.muti_generete(t,A,beam_size=3,device=self.device)
                    res = res.replace('[CLS]','').replace(" ", "")
                    print('res',res)
                    score = myRouge.calc_score([Q],[res])
                    break
            if not found:
                res =self.bert_model.muti_generete(passage,A,beam_size=3,device=self.device)
                res = res.replace('[CLS]','').replace(" ", "")
                print('res',res)
                score = myRouge.calc_score([Q],[res])
            total_score += score
            cnt += 1
        rouge = total_score / cnt
        print('rouge',rouge)
        self.save(self.model_save_path)
            

if __name__ == '__main__':

    trainer = QuestionTrainer()
    train_epoches = 30
    for epoch in range(train_epoches):
        # 训练一个epoch
        trainer.train(epoch)
    

    # train_data, test_data = read_corpus("./round1_train_0907.json", "./roberta_wwm_vocab.txt")
    # print(len(train_data))
    # # print(train_data[:10])
    # vocab_path = "./roberta_wwm_vocab.txt"
    # word2idx = load_chinese_base_vocab(vocab_path=vocab_path, simplfied=True)
    # tokenizer = Tokenizer(word2idx)

    # passage,Q,A =[x[0] for x in train_data],[x[1] for x in train_data],[x[2] for x in train_data]
    # dataset = BertDataset(passage, Q, A, vocab_path)
    # print(len(dataset))
    # dataloader =  DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)
    # print("数据大小为：" + str(len(dataloader)))
    # for token_ids, token_type_ids, target_ids in dataloader:
    #     print(tokenizer.decode(token_ids[0].numpy()))
    #     print(tokenizer.decode(token_ids[1].numpy()))
    #     break


    